#! /usr/bin/env python
# coding=utf-8


# 运行paml的准备工作

# aaml.ctl的模板

# 这里使用的是 Empirical+F 模型


seq_file_name = 'pep_cut_ali.fa'
seq_file_name_aa = 'pep_cut_ali.aa'
tree_file_name = 'spe_tree'
aactl_name = 'aaml.ctl'



from Bio import SeqIO
import os
import shutil
import re
import argparse
import sys




parser = argparse.ArgumentParser(
    description='''为批量运行aaml 进行准备工作 需要输入一个之前cut完 rate4site 生成的路径
    用法:
    conver02_prepare_paml01.py -i OGS_for_testing -n 4 -d jones.dat
    由大天才于2021年11月11日创建于浙江农业大学''')





parser.add_argument('-i',
                help='必须给定，OGS的工作目录')

parser.add_argument('-n',
                help='ncatG,必须给定')

parser.add_argument('-d',
                help='aaml需要的替换模型文件，必须给定')

args = parser.parse_args()

if not args.i or not args.n or not args.d:
    parser.print_help()
    sys.exit()




infile = args.i

ncatG = int(args.n)


dat_file_path  = args.d






paml_tmp = '''

      seqfile = %s
     treefile = %s

      outfile = mlc           * main result file name
        noisy = 3  * 0,1,2,3,9: how much rubbish on the screen
      verbose = 1  * 0: concise; 1: detailed, 2: too much
      runmode = 0  * 0: user tree;  1: semi-automatic;  2: automatic
                   * 3: StepwiseAddition; (4,5):PerturbationNNI; -2: pairwise

      seqtype = 2  * 1:codons; 2:AAs; 3:codons-->AAs
   aaRatefile = my.dat * only used for aa seqs with model=empirical(_F)
                   * dayhoff.dat, jones.dat, wag.dat, mtmam.dat, or your own

        model = 3
                   * models for AAs or codon-translated AAs:
                      * 0:poisson, 1:proportional, 2:Empirical, 3:Empirical+F
                      * 6:FromCodon, 7:AAClasses, 8:REVaa_0, 9:REVaa(nr=189)
        Mgene = 0
                   * AA: 0:rates, 1:separate

        clock = 0  * 0:no clock, 1:global clock; 2:local clock
    fix_alpha = 0  * 0: estimate gamma shape parameter; 1: fix it at alpha
        alpha = 0.5  * initial or fixed alpha, 0:infinity (constant rate)
       Malpha = 0  * different alphas for genes
        ncatG = %s  * # of categories in dG of NSsites models

        getSE = 1  * 0: don't want them, 1: want S.E.s of estimates
 RateAncestor = 1  * (0,1,2): rates (alpha>0) or ancestral states (1 or 2)

   Small_Diff = 1e-6
    cleandata = 0  * remove sites with ambiguity data (1:yes, 0:no)?
*  fix_blength = 0  * 0: ignore, -1: random, 1: initial, 2: fixed
        method = 1   * 0: simultaneous; 1: one branch at a time

''' % (seq_file_name_aa, tree_file_name, ncatG)




for i in os.listdir(infile):
	#print(i)
	target_path = infile+'/'+i+'/'
	# 读取里面的 fasta文件 将其转化为aa文件
	le = 0
	tmp_dic = {}
	seq_nub = 0
	for j in SeqIO.parse(target_path+seq_file_name,'fasta'):
		le = len(j.seq)
		tmp_dic[str(j.name)] = str(j.seq)
		seq_nub+=1
	if le ==0:
		continue
	aa_file = open(target_path+seq_file_name_aa,'w')
	aa_file.write(str(seq_nub)+'  '+str(le)+'\n\n')
	for j in tmp_dic:
		aa_file.write('>\n'+j+'\n'+tmp_dic[j]+'\n\n')

	# 书写aaml.ctl
	aa_ctl = open(target_path+aactl_name,'w')
	aa_ctl.write(paml_tmp)
	aa_ctl.close()
	shutil.copy2(dat_file_path,target_path+'my.dat')
	print(i)
	# 书写


		#print(j.name)
		#print(j.seq)
